from sympy import solve, latex, N
from buss.antlr4_tex2sym.antlr4_tex2sym import tex2sym

def transEqs(expr):
    elist = expr.split('\n')
    begin = 0
    end = len(elist)
    for i in range(end):
        if '\\begin' in elist[i]: begin = i + 1
        if '\\end' in elist[i]: end = i
    if 0 == begin or end == len(elist): return '方程组输入格式错误!'
    elist = elist[begin: end]
    result = []
    for i in elist:
        i = i.strip()
        if i.endswith(r'\\'): i = i[:-2]
        try:
            exp = tex2sym(i)
            if None == exp: return '错误的方程格式: ' + str(i)
            result.append(exp)
        except: return '错误的方程格式: ' + str(i)
    return result

def eq(data):
    expr, unns = data['expr'], data['unns']
    result = {'latex': '', 'text': '', 'float': '', 'error': ''}
    uLen = len(unns)
    if uLen > 1:
        expr = transEqs(expr)
        if type(expr) is str:
            result['error'] = expr
            return result
    elif 1 == uLen:
        unns = unns[0]
        try:
            lexpr = tex2sym(expr)
            if None == lexpr:
                result['error'] = '错误的方程格式: ' + str(expr)
                return result
        except:
            result['error'] = '2错误的方程格式: ' + str(expr)
            return result
    else:
        result['error'] = '未知数输入错误!'
        return result
    try:
        ansr = solve(lexpr, unns)
        if None == ansr:
            result['error'] = '3错误的方程格式: ' + str(lexpr)
            return result
    except:
        result['error'] = '4错误的方程格式: ' + str(lexpr)
        return result
    result['latex'] = latex(ansr)
    result['text'] = str(ansr)
    result['float'] = str([N(i, 2) for i in ansr])
    return result

if __name__ == '__main__':
    a = 'ax^2+bx+c=0'
    b = ['x']
    expr = eq(a, b)
    print(expr)
    a = r'''\left\{ 
\begin{array}{c}
2x+3y+z=1 \\ 
2x+9y+5z=9 \\ 
x+3y+3z=3
\end{array}
\right.'''
    b = ['x', 'y', 'z']
    expr = eq(a, b)
    print(expr)
